package nl.runnable.solr.xslt;

import java.io.IOException;
import java.io.StringWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import javax.xml.transform.Templates;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerException;
import javax.xml.transform.stream.StreamResult;
import javax.xml.transform.stream.StreamSource;

import nl.runnable.solr.util.RequestHandlerHelper;

import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.util.ContentStream;
import org.apache.solr.common.util.ContentStreamBase.StringStream;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.handler.RequestHandlerBase;
import org.apache.solr.handler.XmlUpdateRequestHandler;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.request.SolrRequestHandler;
import org.apache.solr.response.SolrQueryResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * {@link SolrRequestHandler} that transforms incoming XML content using XSLT stylesheets and forwards the request to
 * the {@link XmlUpdateRequestHandler}. The outgoing XML content must be in the XML Update format described in the <a
 * href="http://wiki.apache.org/solr/UpdateXmlMessages">Solr Wiki</a>.
 * 
 * @author Laurens Fridael
 */
public class XsltUpdateHandler extends RequestHandlerBase {

	final Logger logger = LoggerFactory.getLogger(getClass());

	/* Dependencies */

	private TemplatesResolver templatesResolver;

	private RequestHandlerHelper requestHandlerHelper = new RequestHandlerHelper();

	/* Configuration */

	private int scanBufferSize = 4 * 1024;

	private boolean showResultXml = false;

	@Override
	@SuppressWarnings({ "unchecked", "rawtypes" })
	public void init(final NamedList args) {
		super.init(args);
		if (args.get(ConfigurationSettings.SCAN_BUFFER_SIZE) instanceof Integer) {
			setScanBufferSize(((Integer) args.get(ConfigurationSettings.SCAN_BUFFER_SIZE)).intValue());
		}
		setShowResultXml(Boolean.TRUE.equals(args.get(ConfigurationSettings.SHOW_RESULT_XML)));
		final TemplatesResolverFactory templatesResolverFactory = new TemplatesResolverFactory();
		setTemplatesResolver(templatesResolverFactory.createTemplatesResolver(args));
	}

	@Override
	public void handleRequestBody(final SolrQueryRequest request, final SolrQueryResponse response) throws Exception {
		final XmlUpdateRequestHandler requestHandler = getRequestHandlerHelper().getRequestHandler(request,
				XmlUpdateRequestHandler.class);
		if (requestHandler == null) {
			throw new IllegalStateException("Cannot obtain request handler of type " + XmlUpdateRequestHandler.class);
		}
		doHandleRequestBody(request, response, requestHandler);
	}

	protected void doHandleRequestBody(SolrQueryRequest request, final SolrQueryResponse response,
			final RequestHandlerBase followingRequestHandler) throws Exception {
		if (request.getContentStreams() != null) {
			final SolrParams params = request.getParams();
			final Map<String, Object> parameters = getStylesheetParameters(params);

			final List<ContentStream> transformedContentStreams = new ArrayList<ContentStream>();
			for (final ContentStream contentStream : request.getContentStreams()) {
				final ContentStream transformedContentStream = transformContentStream(request, contentStream,
						parameters);
				transformedContentStreams.add(transformedContentStream);
			}
			request = new TransformedSolrQueryRequestWrapper(request, transformedContentStreams);
		}
		followingRequestHandler.handleRequestBody(request, response);
	}

	protected Map<String, Object> getStylesheetParameters(final SolrParams params) {
		final Map<String, Object> parameters = new HashMap<String, Object>();
		for (final Iterator<String> iterator = params.getParameterNamesIterator(); iterator.hasNext();) {
			final String parameterName = iterator.next();
			if (parameterName.startsWith(RequestParameters.PARAM_PREFIX)) {
				final String parameterValue = params.get(parameterName);
				parameters.put(parameterName.substring(RequestParameters.PARAM_PREFIX.length()), parameterValue);
			}
		}
		return parameters;
	}

	protected ContentStream transformContentStream(final SolrQueryRequest request, ContentStream contentStream,
			final Map<String, Object> parameters) throws IOException, TransformerException {
		final PushbackContentStream pushbackContentStream = new PushbackContentStream(contentStream,
				getScanBufferSize());
		final Templates templates = getTemplatesResolver().resolveTemplates(request, pushbackContentStream);
		pushbackContentStream.reset();
		contentStream = pushbackContentStream;
		if (templates != null) {
			final Transformer transformer = templates.newTransformer();
			for (final Entry<String, Object> entry : parameters.entrySet()) {
				transformer.setParameter(entry.getKey(), entry.getValue());
			}
			final StringWriter buffer = new StringWriter();
			transformer.transform(new StreamSource(contentStream.getStream()), new StreamResult(buffer));
			final String resultXml = buffer.toString();
			if (isShowResultXml()) {
				logger.info("Transformation result XML:\n{}", resultXml);
			}
			contentStream = convertToStringStream(contentStream, resultXml);
		}
		return contentStream;
	}

	/**
	 * Converts a given {@link ContentStream} to a {@link StringStream} using the given content.
	 * 
	 * @param contentStream
	 * @param content
	 * @return
	 */
	private StringStream convertToStringStream(final ContentStream contentStream, final String content) {
		final StringStream stringStream = new StringStream(content);
		stringStream.setContentType(contentStream.getContentType());
		stringStream.setSourceInfo(contentStream.getSourceInfo());
		return stringStream;
	}

	/* Dependencies */

	public void setTemplatesResolver(final TemplatesResolver templatesResolver) {
		this.templatesResolver = templatesResolver;
	}

	protected TemplatesResolver getTemplatesResolver() {
		return templatesResolver;
	}

	public void setRequestHandlerHelper(final RequestHandlerHelper requestHandlerHelper) {
		this.requestHandlerHelper = requestHandlerHelper;
	}

	protected RequestHandlerHelper getRequestHandlerHelper() {
		return requestHandlerHelper;
	}

	/* Configuration */

	public void setScanBufferSize(final int bufferSize) {
		this.scanBufferSize = bufferSize;
	}

	protected int getScanBufferSize() {
		return scanBufferSize;
	}

	public void setShowResultXml(final boolean showXml) {
		this.showResultXml = showXml;
	}

	public boolean isShowResultXml() {
		return showResultXml;
	}

	/* TODO: Add MBeanInfo properties */

	@Override
	public String getDescription() {
		return "Transforms incoming XML streams using XSLT to XML Update messages.";
	}

	@Override
	public String getSourceId() {
		// TODO Auto-generated method stub
		return null;
	}

	@Override
	public String getSource() {
		return null;
	}

	@Override
	public String getVersion() {
		// TODO Auto-generated method stub
		return null;
	}

}
